import collections
import os
import numpy
import pandas

from absl import flags

import matplotlib
from matplotlib.lines import Line2D

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()

from coin.strategy.mm.fastfeature.ridge_regress import get_quantile
import coin.pnl.sim_params as sim_params


def divide_chunks(l, n):
  for i in range(0, len(l), n):
    yield l[i:i + n]


def setup_plt():
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['font.family'] = ['Dejavu Sans Mono']
  plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
  plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
  plt.rcParams['font.size'] = 8
  plt.rcParams['legend.fontsize'] = 8
  plt.rcParams['xtick.labelsize'] = 8
  plt.rcParams['ytick.labelsize'] = 8
  plt.rcParams['savefig.facecolor'] = '#ffffff'


def setup_plt_orderbin():
  setup_plt()
  plt.rcParams['figure.figsize'] = 15, 12
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10


def plot_timedret(filename, timedret, suptitle):
  setup_plt()
  plt.rcParams["figure.figsize"] = 8, 4
  timedret['wret'] = timedret['retbps'] * timedret['wgt']
  with PdfPages(filename) as pdf:
    for prodname, productdf in timedret.groupby('product'):
      for category, catedf in productdf.groupby('category'):
        legends = []
        for title, partret in catedf.groupby('title'):
          legends.append(title)
          sumdf = partret.groupby('fwd(ms)').sum()
          plt.plot(sumdf['wret'] / sumdf['wgt'], '.-', rasterized=True, lw=0.5, markersize=1)
          if matplotlib.__version__ >= '3.2.0':
            plt.xscale('symlog', base=2)
          else:
            plt.xscale('symlog', basey=2)
        plt.xlabel('fwd(ms)')
        plt.legend(legends, loc='upper left')
        plt.title(f'{suptitle}_{prodname}_{category}')
        pdf.savefig()
        plt.close()


def plot_orderbin(tssbin,
                  prcbin,
                  sgnbin,
                  dotbin,
                  tagbin=None,
                  start_ts=None,
                  end_ts=None,
                  tag=None):
  if len(sgnbin) == 0:
    return
  idxtag = numpy.array(tagbin) == tag if tagbin is not None and tag is not None else numpy.ones(
      len(sgnbin), dtype=bool)
  for sign, col1, col2, col3 in [(-1, 'tab:red', 'magenta', 'darksalmon'),
                                 (1, 'tab:green', 'limegreen', 'lime')]:
    idxsgn = numpy.array(sgnbin) == sign

    for dottype in ['o', 's', '4', 'x', '|', ',']:
      idxsgnb = numpy.array(dotbin) == dottype
      if idxsgnb.sum() == 0:
        continue
      tssbinab = numpy.array(tssbin)[idxsgn & idxsgnb & idxtag]
      prcbinab = numpy.array(prcbin)[idxsgn & idxsgnb & idxtag]
      col = col1
      mew = None
      if dottype == '|':
        markersize = 1
        mew = 0.5
        col = col2
      elif dottype in ['4', ',']:
        markersize = 1
        mew = 0.5
        col = col3
      elif dottype == 'x':
        markersize = 2
      else:
        markersize = 1
      dti = pandas.DatetimeIndex(tssbinab).to_numpy()
      if start_ts is not None:
        idx = dti.astype(int) >= start_ts
        dti = dti[idx]
        prcbinab = prcbinab[idx]
      if end_ts is not None:
        idx = dti.astype(int) <= end_ts
        dti = dti[idx]
        prcbinab = prcbinab[idx]
      plt.plot(
          dti,
          prcbinab,
          markersize=markersize,
          mew=mew,
          color=col,
          ls='',
          marker=dottype)
    plt.plot(
        pandas.DatetimeIndex(numpy.array(tssbin)[idxsgn & idxtag]).to_numpy(),
        numpy.array(prcbin)[idxsgn & idxtag],
        markersize=markersize,
        lw=0.2,
        alpha=0.5,
        color=col,
        ls='-',
        marker='')

    legend_elements = [
        Line2D([0], [0],
               marker='o', color=col1, label='MakerFill', markerfacecolor='g', markersize=2),
        Line2D([0], [0],
               marker='s', color=col1, label='TakerFill', markerfacecolor='g', markersize=2),
        Line2D([0], [0],
               marker='x',
               color=col1,
               label='Error/Reject',
               markerfacecolor='g',
               markersize=4)
    ]
    plt.legend(handles=legend_elements, loc="upper right")


def gen_miss_fill_lines_from_oedf(oedf):
  subidx = oedf['type'] == 'ORDER_SUBMITTED'
  canidx = oedf['type'] == 'CANCEL_CONFIRMED'
  fillidx = oedf['type'] == 'ORDER_FILLED'
  subdf = oedf.loc[subidx]
  if 'internal_order_id' in subdf:
    oedf = oedf.merge(
        subdf[['internal_order_id', 'timestamp']],
        on='internal_order_id', how='left', suffixes=('', '_submit'))
  else:
    oedf = oedf.merge(
        subdf[['order_id', 'timestamp']],
        on='order_id', how='left', suffixes=('', '_submit'))
  candf = oedf.loc[canidx]
  filldf = oedf.loc[fillidx]

  for lstyle, domain_df, column in [(':', candf, 'order_price'), ('-', filldf, 'fill_price')]:
    for sign, targetdf in domain_df.groupby('sign'):
      tssub = targetdf['timestamp_submit']
      tscol = targetdf['timestamp']
      prcsub = targetdf['order_price']
      prccol = targetdf[column]
      if int(sign) == 1:
        col = 'green'
      elif int(sign) == -1:
        col = 'red'
      else:
        continue

      tsss = pandas.DataFrame(zip(tssub, tscol, tscol * numpy.nan)).to_numpy().flatten()
      prcs = pandas.DataFrame(zip(prcsub, prccol, prccol * numpy.nan)).to_numpy().flatten()
      yield tsss, prcs, lstyle, col


def plot_miss_fill_lines_from_oedf(oedf):
  for tss, prcs, linestyle, linecol in gen_miss_fill_lines_from_oedf(oedf):
    plt.plot(pandas.DatetimeIndex(tss), prcs, ls=linestyle, color=linecol, lw=0.2)


def plot_orderdf_plotonly(stratname, simstat):
  if simstat is None:
    return
  combdf = simstat['combdf']
  orderdf = simstat['orderdf']
  prcdict = simstat['prcdict']
  setup_plt()
  plt.rcParams["figure.figsize"] = 8, 12
  plt.subplot(411)
  # plt.plot(pandas.DatetimeIndex(combdf['timestamp']), combdf['pnl'], lw=0.5)
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']), combdf['pnl_net'], lw=0.5)
  if combdf['fill_qty'].sum() > 0:
    net_ret = combdf['pnl_net'].iloc[-1] / (combdf['fill_qty']).sum()
  else:
    net_ret = 0
  plt.title("%s pnl. net ret/bps: %.2f" % (stratname, net_ret * 1e4))
  # TODO: vc = volume_overlay3(ax2, quotes)
  # ax2.add_collection(vc)

  plt.subplot(412)
  plt.tick_params(axis='y', which='both', labelleft='off', labelright='on')
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']), combdf['open_pos'], lw=0.5)
  plt.title("(sim) pos & price")
  plt.ylabel("position")

  plt.twinx()
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']),
           combdf['bid0_qty'],
           color='peachpuff',
           ls='--',
           lw=0.5)
  invask0qty = combdf['ask0_qty'].max() + combdf['bid0_qty'].max() - combdf['ask0_qty']
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']),
           invask0qty,
           color='powderblue',
           ls='--',
           lw=0.5)
  plt.axis('off')

  plt.twinx()
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']),
           combdf['bid0_price'].fillna(method='ffill'),
           'g:',
           pandas.DatetimeIndex(combdf['timestamp']),
           combdf['ask0_price'].fillna(method='ffill'),
           'r:',
           lw=0.2,
           drawstyle='steps-post')
  plt.ylabel("price")

  plot_orderbin(**prcdict)

  plt.subplot(413)
  plt.axis('off')
  plt.subplot(4, 4, 13)

  taker_qty_sum = (orderdf['fill_qty'] * (orderdf['fill_type'] == 'TAKER')).sum()
  maker_qty_sum = (orderdf['fill_qty'] * (orderdf['fill_type'] == 'MAKER')).sum()

  fill_type = 'TAKER' if taker_qty_sum > maker_qty_sum else 'MAKER'
  taker_or_maker_df = orderdf[orderdf['fill_type'] == fill_type].reset_index(drop=True)
  takqty, fret = get_quantile(taker_or_maker_df['fill_qty'], taker_or_maker_df['fill_ret'], 10)
  plt.plot(takqty, fret, 'g.-', markersize=1, linewidth=0.5)
  plt.xlabel("%s qty" % fill_type)
  plt.ylabel("%s ret" % fill_type)
  plt.title("fill qty vs fill return")

  filldf = orderdf.loc[orderdf['fill_qty'] > 0]

  plt.subplot(4, 4, 14)
  buyfill = filldf.loc[filldf['sign'] == 1]
  sellfill = filldf.loc[filldf['sign'] == -1]
  if len(buyfill) > 0 and len(sellfill) > 0:
    buypdiff = buyfill['bid0_price_-1s_fwd'] - buyfill['fill_price']
    buyqt, bret = get_quantile(buypdiff, buyfill['fill_ret'], 10)
    sellpdiff = -(sellfill['ask0_price_-1s_fwd'] - sellfill['fill_price'])
    sellqt, sret = get_quantile(sellpdiff, sellfill['fill_ret'], 10)
    plt.plot(buyqt, bret, 'g.-', sellqt, sret, 'r.-', markersize=1, linewidth=0.5)

  plt.subplot(4, 4, 15)
  if len(filldf) > 0:
    elapsedt, fret = get_quantile(filldf['elapsec'], filldf['fill_ret'], 10)
    plt.plot(elapsedt, fret, 'g.-', markersize=1, linewidth=0.5)
    plt.xlabel("elapsed time")
    plt.ylabel("fill ret")

  candf = orderdf.loc[(orderdf['type'] == 'CANCELLED')]
  elapsedt, qposrat = get_quantile(candf['elapsec'], candf['queue_pos_ratio'], 10)
  plt.title('order elapsed time vs (return & qpos)')
  plt.twinx()
  plt.plot(numpy.nan, numpy.nan, 'g.-', markersize=1, linewidth=0.5)
  plt.plot(elapsedt, qposrat, 'r.-', markersize=1, linewidth=0.5)
  plt.ylabel("qpos rat")
  plt.legend(['fillret', 'qpos'], loc='upper right')

  plt.subplot(4, 4, 16)

  if len(filldf) > 0:
    nfilldf = filldf.loc[~numpy.isnan(filldf['queue_pos'])]
    fillpos, fillret = get_quantile(nfilldf['queue_pos'], nfilldf['fill_ret'], 10)
    plt.plot(fillpos, fillret, 'g-', markersize=1, linewidth=0.5)
    plt.title('queue_pos vs fill_ret')
    plt.xlabel('queue pos before fill')
    plt.tick_params(axis='y', which='both', labelleft='off', labelright='on')


def plot_orderdf(root_dir, filename_str, stratname, simstat):
  if simstat is None:
    return
  plot_orderdf_plotonly(stratname, simstat)
  try:
    # plt.tight_layout()
    plt.savefig(os.path.join(root_dir, "%s.png" % filename_str))
  except Exception:
    pass
  finally:
    plt.close()


def plot_detail_plotonly(stratname, simstat):
  if simstat is None:
    return
  setup_plt_orderbin()
  combdf = simstat['combdf']
  plt.tick_params(axis='y', which='both', labelleft='off', labelright='on')
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']),
           combdf['open_pos'],
           lw=0.5,
           drawstyle='steps-post')
  plt.title(f"(sim) pos & price {stratname}")
  plt.ylabel("position")

  plt.twinx()
  plt.plot(pandas.DatetimeIndex(combdf['timestamp']),
           combdf['pnl_net'],
           'k:',
           lw=0.2,
           markersize=0.2,
           drawstyle='steps-post')
  plt.axis('off')
  plt.legend(['pnl'], loc='upper left')

  if 'bbodf' in simstat:
    plt.twinx()
    bbodf = simstat['bbodf']
    bbodf = bbodf.loc[~numpy.isnan(bbodf['bid0_price'])]
    bbodf = bbodf.loc[~numpy.isnan(bbodf['ask0_price'])]
    plt.plot(pandas.DatetimeIndex(bbodf['timestamp']),
             bbodf['bid0_price'].fillna(method='ffill'),
             'g:',
             pandas.DatetimeIndex(bbodf['timestamp']),
             bbodf['ask0_price'].fillna(method='ffill'),
             'r:',
             lw=0.2,
             drawstyle='steps-post')
    plt.ylabel("price")
  plot_orderbin(**simstat['prcdict'])
  plot_miss_fill_lines_from_oedf(simstat['orderdf'])


def plot_detail(root_dir, filename_str, stratname, simstat):
  if simstat is None:
    return
  plot_detail_plotonly(stratname, simstat)

  try:
    plt.savefig(os.path.join(root_dir, "%s_detail.png" % filename_str), dpi=300)
  except Exception:
    pass
  finally:
    plt.close()


def merge_group_stat(group_stats):
  df = pandas.DataFrame(group_stats)
  grp_series = df.sum()
  grp_ids = [column.replace("_acc_qty", "") for column in df.keys() if column.find("acc_qty") >= 0]

  visual_stat = collections.OrderedDict()
  for grp_id in grp_ids:
    for tm in ['taker', 'maker']:
      visual_stat[f'{grp_id}_{tm}_ret'] = grp_series[f'{grp_id}_{tm}_pnl'] / \
          grp_series[f'{grp_id}_{tm}_qty']
      visual_stat[f'{grp_id}_{tm}_qty'] = grp_series[f'{grp_id}_{tm}_qty']
      visual_stat[f'{grp_id}_{tm}_qty_ratio'] = grp_series[f'{grp_id}_{tm}_qty'] / \
        grp_series[f'{grp_id}_acc_qty']
  return visual_stat


def plot_markreturn(root_dir, filename_str, first_order_df):
  def plot_edge_fr_ret_ev(edge, fr, ret):
    plt.plot(edge, fr, 'g.-', numpy.nan, numpy.nan, 'r.-', lw=0.5, markersize=1)
    plt.legend(['fr', 'ret'], loc='upper left')
    plt.twinx()
    plt.plot(edge, ret, 'r.-', lw=0.5, markersize=1)

  def smallplot():
    plt.rcParams['figure.figsize'] = 6, 6
    plt.rcParams['font.family'] = ['Dejavu Sans Mono']
    plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
    plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
    plt.rcParams['font.size'] = 10
    plt.rcParams['legend.fontsize'] = 8
    plt.rcParams['xtick.labelsize'] = 8
    plt.rcParams['ytick.labelsize'] = 8

  smallplot()
  for maktak, splt in [('taker', 211), ('maker', 212)]:
    plt.subplot(splt)
    seldf = first_order_df.loc[first_order_df[f'is_{maktak}']]

    nbucket = 20
    edges = seldf['payup_bp']
    fills = (seldf['fill_qty'] > 0)
    rets = (seldf['fill_qty'] > 0) * (seldf['orderp_ret_120s'] - seldf[f'{maktak}_fee'])
    _, cols = zip(*sorted([(float(orcol.split('_')[-1].replace('s', '')), orcol)
                           for orcol in seldf.columns
                           if orcol.startswith("orderp_ret")]))
    print(maktak)
    print(seldf[list(cols)][seldf['fill_qty'] > 0].mean())
    plt.title("%s frmean: %.2f avgpup: %.2f avgrt: %.2f" % (maktak, (seldf['fill_qty'] > 0).mean(),
                                                            seldf['payup_bp'].mean(),
                                                            rets.sum() / fills.sum() * 1e4))
    eq, fq = get_quantile(edges, fills, nbucket)
    eq, retq = get_quantile(edges, rets, nbucket)
    retq = pandas.Series(numpy.array(retq) / numpy.array(fq))
    retq.fillna(0)
    plot_edge_fr_ret_ev(eq, fq, retq)
    plt.axhline(y=0, lw=1)
  plt.tight_layout()
  plt.savefig(os.path.join(root_dir, "%s_mark_return.png" % filename_str))
  plt.close()


def plot_multidays(groupkey, groupval, subdf, sim_stats, idxs):
  setup_plt()
  plt.rcParams["figure.figsize"] = 14, 12

  first_order_dfs = []
  pnldfs = []
  group_stats_sels = []
  max_tps = 0

  for i in idxs:
    sim_stat = sim_stats[i]
    pnldf = sim_stat[2]['combdf']
    group_stat = sim_stat[2]['group_stat']
    first_order_df = sim_stat[2]['first_order_df']
    pnldf['pnl_diff'] = pnldf['pnl_net'].diff().fillna(0)
    pnldfs.append(pnldf)
    first_order_dfs.append(first_order_df)
    group_stats_sels.append(group_stat)
    max_tps = max(max_tps, group_stat['tps'])

  visual_stat = merge_group_stat(group_stats_sels)

  wholepnldf = pandas.concat(pnldfs)
  wholepnldf = wholepnldf.sort_values('timestamp').reset_index(drop=True)
  wholepnldf['pnl_net'] = wholepnldf['pnl_diff'].cumsum()

  dfkeyname = "__".join(
      [f"{sim_params.get_abstract_name(key)}:{str(val)}" for key, val in zip(groupkey, groupval)])
  print(dfkeyname)
  mindate, maxdate = subdf['trading_date'].min(), subdf['trading_date'].max()
  title = f'{mindate}-{maxdate}_{dfkeyname}'
  filename = f'{flags.FLAGS.sim_root}/{title}.png'

  shortten = lambda key: key.replace("TRADE_FEED", "T").replace("BOOK_FEED", "B").replace(
      "taker", "tak").replace("maker", "mak")
  shortten2 = lambda key: key.replace("PERPETUAL", "SWAP").replace("QUARTER", "QT").replace(
      "THIS_WEEK", "TW").replace("NEXT_WEEK", "NW")
  shortten3 = lambda key: key.replace("Futures", "Fut").replace('qty', 'q').replace('ratio', 'rat')

  def format_with_key(key, value):
    if numpy.isnan(value):
      return 'N/A'
    if key.endswith('qty'):
      return format(value, '.3f')
    else:
      return format(value * 1e4, '.1f') + 'bp'

  outtuple = [
      "%s: %s" % (shortten3(shortten2(shortten(key))), format_with_key(key, value)) for key,
      value in visual_stat.items()
  ]

  plt.rcParams["figure.figsize"] = 12, 16
  setup_plt()
  df = pandas.concat([pandas.DataFrame(a).T for a in divide_chunks(outtuple, 3)])
  plt.subplot(311)
  plt.plot(pandas.DatetimeIndex(wholepnldf['timestamp']), wholepnldf['pnl_net'])
  plt.title(title)

  plt.subplot(312)
  plt.plot(pandas.DatetimeIndex(wholepnldf['timestamp']), wholepnldf['open_pos'])

  plt.subplot(313)
  plt.title("max_tps: %s" % max_tps)
  tbl = plt.table(cellText=df.values, loc='center')
  tbl.scale(1, 1.5)
  tbl.auto_set_font_size(False)
  tbl.set_fontsize(9)
  plt.axis('off')

  plt.savefig(filename)
  plt.close()

  first_order_df = pandas.concat(first_order_dfs)
  plot_markreturn(flags.FLAGS.sim_root, title, first_order_df)
